from PyQt5.QtWidgets import (QMainWindow, QTabWidget, QWidget, QVBoxLayout,
                             QHBoxLayout, QPushButton, QFileDialog, QLabel,
                             QComboBox, QFormLayout, QSpinBox, QDateEdit,
                             QApplication, QDoubleSpinBox, QGroupBox, QProgressDialog,
                             QProgressBar,
                             )
from PyQt5.QtCore import Qt, QDate, QSettings, QSize, QPoint, QTimer
from fontTools.misc.plistlib import end_date
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from core.analyzer import DowTheoryAnalyzer
from core.backtester import Backtester
from data.loader import DataLoader
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.dates import DateFormatter, date2num
from datetime import datetime, timedelta
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
import pandas as pd
from gui.charts import setup_chinese_fonts

# 确保字体设置正确加载
setup_chinese_fonts()


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("道氏理论分析系统")
        self.setGeometry(100, 100, 1200, 800)
        self.start_date = None
        self.end_date = None

        # 初始化组件
        self.analyzer = None
        self.backtester = Backtester()
        self.current_data = None

        # 创建主界面
        self._create_ui()

        # 加载设置
        self.settings = QSettings("YourCompany", "DowTheoryAnalyzer")
        self._load_settings()

        # 创建状态栏进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setMaximumWidth(200)
        self.progress_bar.setVisible(False)
        self.statusBar().addPermanentWidget(self.progress_bar)

    def _create_ui(self):
        """创建界面布局"""
        main_widget = QWidget()
        self.setCentralWidget(main_widget)

        # 主布局
        layout = QVBoxLayout()
        main_widget.setLayout(layout)

        # 控制面板
        control_panel = self._create_control_panel()
        layout.addLayout(control_panel)

        # 标签页
        self.tabs = QTabWidget()
        layout.addWidget(self.tabs)

        # 添加标签页
        self._create_analysis_tab()
        self._create_backtest_tab()

    def _load_settings(self):
        """加载保存的设置"""
        # 窗口大小和位置
        self.resize(self.settings.value("window_size", QSize(1200, 800)))
        self.move(self.settings.value("window_position", QPoint(100, 100)))

        # 分析参数
        self.param_primary.setValue(int(self.settings.value("param_primary", 30)))
        self.param_secondary.setValue(int(self.settings.value("param_secondary", 10)))
        self.volume_threshold.setValue(float(self.settings.value("volume_threshold", 1.2)))

    def _save_settings(self):
        """保存当前设置"""
        # 窗口大小和位置
        self.settings.setValue("window_size", self.size())
        self.settings.setValue("window_position", self.pos())

        # 分析参数
        self.settings.setValue("param_primary", self.param_primary.value())
        self.settings.setValue("param_secondary", self.param_secondary.value())
        self.settings.setValue("volume_threshold", self.volume_threshold.value())

        # 时间范围
        if hasattr(self, 'start_date') and self.start_date:
            self.settings.setValue("start_date", self.start_date.isoformat())
        if hasattr(self, 'end_date') and self.end_date:
            self.settings.setValue("end_date", self.end_date.isoformat())

    def closeEvent(self, event):
        """窗口关闭时保存设置"""
        self._save_settings()
        super().closeEvent(event)

    def _create_control_panel(self) -> QHBoxLayout:
        """创建顶部控制面板"""
        main_panel = QHBoxLayout()

        # 文件加载区域
        file_group = QGroupBox("数据管理")
        file_layout = QVBoxLayout()
        btn_load = QPushButton("加载数据")
        btn_load.clicked.connect(self._load_data)
        file_layout.addWidget(btn_load)
        file_group.setLayout(file_layout)
        main_panel.addWidget(file_group)

        # 时间控制区域
        time_group = QGroupBox("时间控制")
        time_layout = QVBoxLayout()

        # 时间范围选择
        range_layout = QHBoxLayout()
        range_layout.addWidget(QLabel("开始日期:"))
        self.date_start = QDateEdit()
        self.date_start.setCalendarPopup(True)
        self.date_start.setDisplayFormat("yyyy-MM-dd")
        range_layout.addWidget(self.date_start)

        range_layout.addWidget(QLabel("结束日期:"))
        self.date_end = QDateEdit()
        self.date_end.setCalendarPopup(True)
        self.date_end.setDisplayFormat("yyyy-MM-dd")
        range_layout.addWidget(self.date_end)

        time_layout.addLayout(range_layout)

        # 预设时间范围按钮
        preset_layout = QHBoxLayout()
        btn_1m = QPushButton("1个月")
        btn_1m.clicked.connect(lambda: self._set_preset_range(30))
        btn_3m = QPushButton("3个月")
        btn_3m.clicked.connect(lambda: self._set_preset_range(90))
        btn_6m = QPushButton("6个月")
        btn_6m.clicked.connect(lambda: self._set_preset_range(180))
        btn_1y = QPushButton("1年")
        btn_1y.clicked.connect(lambda: self._set_preset_range(365))
        btn_all = QPushButton("全部")
        btn_all.clicked.connect(self._set_full_range)

        preset_layout.addWidget(btn_1m)
        preset_layout.addWidget(btn_3m)
        preset_layout.addWidget(btn_6m)
        preset_layout.addWidget(btn_1y)
        preset_layout.addWidget(btn_all)
        time_layout.addLayout(preset_layout)

        # 导航按钮
        nav_layout = QHBoxLayout()
        btn_prev = QPushButton("←")
        btn_prev.clicked.connect(self._navigate_prev)
        btn_apply = QPushButton("应用范围")
        btn_apply.clicked.connect(self._apply_date_range)
        btn_next = QPushButton("→")
        btn_next.clicked.connect(self._navigate_next)

        nav_layout.addWidget(btn_prev)
        nav_layout.addWidget(btn_apply)
        nav_layout.addWidget(btn_next)
        time_layout.addLayout(nav_layout)

        time_group.setLayout(time_layout)
        main_panel.addWidget(time_group)

        # 分析参数区域
        param_group = QGroupBox("分析参数")
        param_layout = QFormLayout()

        self.param_primary = QSpinBox()
        self.param_primary.setRange(10, 100)
        self.param_primary.setValue(30)
        param_layout.addRow("主要周期:", self.param_primary)

        self.param_secondary = QSpinBox()
        self.param_secondary.setRange(5, 50)
        self.param_secondary.setValue(10)
        param_layout.addRow("次要周期:", self.param_secondary)

        self.volume_threshold = QDoubleSpinBox()
        self.volume_threshold.setRange(1.0, 3.0)
        self.volume_threshold.setSingleStep(0.1)
        self.volume_threshold.setValue(1.2)
        param_layout.addRow("成交量阈值:", self.volume_threshold)

        param_group.setLayout(param_layout)
        main_panel.addWidget(param_group)

        # 操作按钮区域
        action_group = QGroupBox("操作")
        action_layout = QVBoxLayout()

        btn_analyze = QPushButton("执行分析")
        btn_analyze.clicked.connect(self._run_analysis)
        action_layout.addWidget(btn_analyze)

        btn_backtest = QPushButton("执行回测")
        btn_backtest.clicked.connect(self._run_backtest)
        action_layout.addWidget(btn_backtest)

        action_group.setLayout(action_layout)
        main_panel.addWidget(action_group)

        return main_panel

    def _create_analysis_tab(self):
        """创建分析结果标签页"""
        tab = QWidget()
        layout = QVBoxLayout()
        tab.setLayout(layout)

        # 价格图表和工具栏
        price_box = QWidget()
        price_layout = QVBoxLayout()
        price_box.setLayout(price_layout)

        self.price_chart = MplCanvas()
        price_layout.addWidget(self.price_chart)

        # 添加Matplotlib工具栏
        self.price_toolbar = NavigationToolbar(self.price_chart, self)
        price_layout.addWidget(self.price_toolbar)

        layout.addWidget(price_box)

        self.tabs.addTab(tab, "趋势分析")

    def _create_backtest_tab(self):
        """创建回测标签页"""
        tab = QWidget()
        layout = QVBoxLayout()
        tab.setLayout(layout)

        # 资金曲线图表和工具栏
        equity_box = QWidget()
        equity_layout = QVBoxLayout()
        equity_box.setLayout(equity_layout)

        self.equity_chart = MplCanvas()
        equity_layout.addWidget(self.equity_chart)

        self.equity_toolbar = NavigationToolbar(self.equity_chart, self)
        equity_layout.addWidget(self.equity_toolbar)

        layout.addWidget(equity_box)

        # 回测结果图表
        self.equity_chart = MplCanvas()
        layout.addWidget(self.equity_chart)

        # 绩效指标
        self.metrics_label = QLabel("回测结果将在此显示...")
        self.metrics_label.setAlignment(Qt.AlignLeft)
        layout.addWidget(self.metrics_label)

        # 回测按钮
        btn_backtest = QPushButton("执行回测")
        btn_backtest.clicked.connect(self._run_backtest)
        layout.addWidget(btn_backtest)

        self.tabs.addTab(tab, "策略回测")

    def _load_data(self):
        """加载数据文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择数据文件", "",
            "CSV文件 (*.csv);;Excel文件 (*.xls *.xlsx)"
        )
        if file_path:
            try:
                loader = DataLoader()
                self.current_data = loader.load(file_path)

                # 设置日期选择器的范围
                min_date = self.current_data.index.min()
                max_date = self.current_data.index.max()

                # 加载保存的时间范围或使用默认范围
                saved_start = self.settings.value("start_date")
                saved_end = self.settings.value("end_date")

                if saved_start and saved_end:
                    try:
                        start_date = pd.Timestamp(saved_start)
                        end_date = pd.Timestamp(saved_end)

                        # 确保保存的时间范围在数据范围内
                        if start_date < min_date:
                            start_date = min_date
                        if end_date > max_date:
                            end_date = max_date
                        if start_date > end_date:
                            start_date = min_date
                            end_date = max_date

                        self.start_date = start_date
                        self.end_date = end_date
                    except:
                        self.start_date = min_date
                        self.end_date = max_date
                else:
                    self.start_date = min_date
                    self.end_date = max_date

                self.date_start.setDate(QDate(self.start_date.year, self.start_date.month, self.start_date.day))
                self.date_end.setDate(QDate(self.end_date.year, self.end_date.month, self.end_date.day))

                # 默认显示完整时间范围
                self.start_date = min_date
                self.end_date = max_date

                self.statusBar().showMessage(f"成功加载: {file_path}，共 {len(self.current_data)} 条数据", 5000)
            except Exception as e:
                self.statusBar().showMessage(f"加载失败: {str(e)}", 5000)

    def _apply_date_range(self):
        """应用选择的时间范围"""
        if self.current_data is None:
            self.statusBar().showMessage("请先加载数据!", 3000)
            return

        # 获取选择的日期
        start_date = self.date_start.date().toPyDate()
        end_date = self.date_end.date().toPyDate()

        # 将日期转换为 Pandas 时间戳格式
        self.start_date = pd.Timestamp(start_date)
        self.end_date = pd.Timestamp(end_date)

        # 验证日期范围
        data_min = self.current_data.index.min()
        data_max = self.current_data.index.max()

        if self.start_date < data_min or self.end_date > data_max:
            self.statusBar().showMessage("日期超出数据范围!", 3000)
            return

        if self.start_date > self.end_date:
            self.statusBar().showMessage("开始日期不能晚于结束日期!", 3000)
            return

        # 如果有分析结果，更新图表
        if hasattr(self, 'analyzer') and self.analyzer is not None:
            self._update_charts(self.analyzer)
            self.statusBar().showMessage(f"已更新时间范围: {start_date} 至 {end_date}", 3000)

    def _set_preset_range(self, days):
        """设置预设时间范围"""
        if self.current_data is None:
            self.statusBar().showMessage("请先加载数据!", 3000)
            return

        end_date = self.end_date if self.end_date else self.current_data.index.max()
        start_date = end_date - pd.Timedelta(days=days)

        # 确保时间范围在数据范围内
        min_date = self.current_data.index.min()
        if start_date < min_date:
            start_date = min_date

        self.start_date = start_date
        self.end_date = end_date

        self.date_start.setDate(QDate(start_date.year, start_date.month, start_date.day))
        self.date_end.setDate(QDate(end_date.year, end_date.month, end_date.day))

        self._apply_date_range()

    def _set_full_range(self):
        """设置完整时间范围"""
        if self.current_data is None:
            self.statusBar().showMessage("请先加载数据!", 3000)
            return

        min_date = self.current_data.index.min()
        max_date = self.current_data.index.max()

        self.start_date = min_date
        self.end_date = max_date

        self.date_start.setDate(QDate(min_date.year, min_date.month, min_date.day))
        self.date_end.setDate(QDate(max_date.year, max_date.month, max_date.day))

        self._apply_date_range()

    def _navigate_prev(self):
        """导航到前一个时间段"""
        if self.current_data is None or not hasattr(self, 'start_date') or not self.start_date:
            self.statusBar().showMessage("请先加载数据并设置时间范围!", 3000)
            return

        # 计算当前时间范围长度
        delta = self.end_date - self.start_date

        # 计算新的时间范围
        new_end = self.start_date - pd.Timedelta(days=1)
        new_start = new_end - delta

        # 确保时间范围在数据范围内
        min_date = self.current_data.index.min()
        if new_start < min_date:
            new_start = min_date
            new_end = new_start + delta

        self.start_date = new_start
        self.end_date = new_end

        self.date_start.setDate(QDate(new_start.year, new_start.month, new_start.day))
        self.date_end.setDate(QDate(new_end.year, new_end.month, new_end.day))

        self._apply_date_range()

    def _navigate_next(self):
        """导航到下一个时间段"""
        if self.current_data is None or not hasattr(self, 'start_date') or not self.start_date:
            self.statusBar().showMessage("请先加载数据并设置时间范围!", 3000)
            return

        # 计算当前时间范围长度
        delta = self.end_date - self.start_date

        # 计算新的时间范围
        new_start = self.end_date + pd.Timedelta(days=1)
        new_end = new_start + delta

        # 确保时间范围在数据范围内
        max_date = self.current_data.index.max()
        if new_end > max_date:
            new_end = max_date
            new_start = new_end - delta
            if new_start < self.start_date:  # 防止无限循环
                new_start = self.start_date + pd.Timedelta(days=1)

        self.start_date = new_start
        self.end_date = new_end

        self.date_start.setDate(QDate(new_start.year, new_start.month, new_start.day))
        self.date_end.setDate(QDate(new_end.year, new_end.month, new_end.day))

        self._apply_date_range()

    def _run_analysis(self):
        """执行分析（带性能优化）"""
        if self.current_data is None:
            self.statusBar().showMessage("请先加载数据!", 3000)
            return

        try:
            # 显示进度条
            self.progress_bar.setVisible(True)
            self.progress_bar.setRange(0, 100)
            self.progress_bar.setValue(0)
            self.statusBar().showMessage("开始分析数据...")
            QApplication.processEvents()

            # 显示等待光标
            QApplication.setOverrideCursor(Qt.WaitCursor)
            QApplication.processEvents()  # 确保UI更新

            # 应用当前选择的时间范围
            if self.start_date and self.end_date:
                filtered_data = self.current_data[
                    (self.current_data.index >= self.start_date) &
                    (self.current_data.index <= self.end_date)
                    ]
            else:
                filtered_data = self.current_data.copy()

            if len(filtered_data) < 30:
                self.statusBar().showMessage("数据量不足，请选择更长时间范围!", 5000)
                QApplication.restoreOverrideCursor()
                return

            # 更新进度
            self.progress_bar.setValue(30)
            self.statusBar().showMessage("识别关键点位...")
            QApplication.processEvents()

            # 使用进度条
            progress = QProgressDialog("正在分析数据...", "取消", 0, 100, self)
            progress.setWindowModality(Qt.WindowModal)
            progress.setMinimumDuration(0)
            progress.setValue(10)
            QApplication.processEvents()

            self.analyzer = DowTheoryAnalyzer(
                window_primary=self.param_primary.value(),
                window_secondary=self.param_secondary.value(),
                volume_threshold=self.volume_threshold.value()
            )
            progress.setValue(40)
            QApplication.processEvents()

            result = self.analyzer.analyze(filtered_data)
            progress.setValue(70)
            QApplication.processEvents()

            self.progress_bar.setValue(70)
            self.statusBar().showMessage("生成信号...")
            QApplication.processEvents()

            self.progress_bar.setValue(100)
            self.statusBar().showMessage("分析完成!", 3000)

            # 更新图表
            self._update_charts(result)
            progress.setValue(100)

            self.statusBar().showMessage("分析完成!", 3000)
        except Exception as e:
            self.statusBar().showMessage(f"分析失败: {str(e)}", 5000)
        finally:
            # 恢复光标
            QApplication.restoreOverrideCursor()
            if 'progress' in locals():
                progress.close()
            QTimer.singleShot(3000, lambda: self.progress_bar.setVisible(False))

    def _run_backtest(self):
        """执行回测"""
        if not hasattr(self.analyzer, 'signals'):
            self.statusBar().showMessage("请先完成分析!", 3000)
            return

        try:
            # 直接使用分析器中的数据（已筛选过时间范围）
            filtered_data = self.analyzer.data
            filtered_signals = self.analyzer.signals['Signal']

            # 将日期转换为pandas.Timestamp 类型
            start_ts = pd.Timestamp(self.start_date) if self.start_date else None
            end_ts = pd.Timestamp(self.end_date) if self.end_date else None

            results = self.backtester.run_backtest(
                data=filtered_data,  # 使用分析器中的筛选后数据
                signals=filtered_signals,  # 使用分析器中的筛选后信号
                start_date=start_ts,  # 不再需要额外筛选
                end_date=end_ts
            )

            # 显示回测结果
            self._plot_equity_curve(results['positions'])
            self._show_metrics(results['performance'])
            self.statusBar().showMessage("回测完成!", 3000)
        except Exception as e:
            self.statusBar().showMessage(f"回测失败: {str(e)}", 5000)

    def _update_charts(self, result):
        """更新分析图表，使用当前时间范围"""
        try:
            # 直接使用分析器数据
            data = self.analyzer.data
            signals = self.analyzer.signals

            # 显示等待光标
            QApplication.setOverrideCursor(Qt.WaitCursor)
            QApplication.processEvents()

            # 调整子图高度比例（价格:成交量 = 3:1）
            self.price_chart.figure.clear()
            gs = self.price_chart.figure.add_gridspec(2, 1, height_ratios=[3, 1])
            ax1 = self.price_chart.figure.add_subplot(gs[0])
            ax2 = self.price_chart.figure.add_subplot(gs[1])

            # 提取当前时间范围的数据
            if self.start_date and self.end_date:
                data = result.data[
                    (result.data.index >= self.start_date) &
                    (result.data.index <= self.end_date)
                    ]
            else:
                data = result.data

            if data.empty:
                self.statusBar().showMessage("选定时间范围无数据!", 3000)
                return

            # 绘制价格线
            ax1.plot(data.index, data['Close'],
                    label='Price', color='black', linewidth=1.5)

            # 绘制趋势线
            if result.trend_lines['primary_up']:
                line = result.trend_lines['primary_up']
                # 只绘制在可见范围内的趋势线
                start_visible = max(line['start_date'], data.index.min())
                end_visible = min(line['end_date'], data.index.max())


                if start_visible < end_visible:
                    dates = [start_visible, end_visible]
                    start_ordinal = start_visible.toordinal()
                    end_ordinal = end_visible.toordinal()
                    values = [
                        line['slope'] * start_ordinal + line['intercept'],
                        line['slope'] * end_ordinal + line['intercept']
                    ]
                    ax1.plot(dates, values, 'g--', linewidth=1.5, label='Primary Up Trend')

            if result.trend_lines['primary_down']:
                line = result.trend_lines['primary_down']
                # 只绘制在可见范围内的趋势线
                start_visible = max(line['start_date'], data.index.min())
                end_visible = min(line['end_date'], data.index.max())

                if start_visible < end_visible:
                    dates = [start_visible, end_visible]
                    values = [
                        line['slope'] * start_visible.toordinal() + line['intercept'],
                        line['slope'] * end_visible.toordinal() + line['intercept']
                    ]
                    ax1.plot(dates, values, 'r--', linewidth=1.5, label='Primary Down Trend')

            # 绘制买卖信号
            signals = result.signals.loc[data.index]
            if not signals.empty:
                buy_signals = signals[signals['Signal'] == 1]
                sell_signals = signals[signals['Signal'] == -1]

                if not buy_signals.empty:
                    ax1.scatter(buy_signals.index,
                                data.loc[buy_signals.index, 'Close'] * 0.98,
                                color='blue', s=100, marker='^', label='Buy Signal')

                if not sell_signals.empty:
                    ax1.scatter(sell_signals.index,
                                data.loc[sell_signals.index, 'Close'] * 1.02,
                                color='red', s=100, marker='v', label='Sell Signal')

            # 设置价格图表格式
            ax1.set_title(f'道氏理论分析 ({self.start_date} 至 {self.end_date})')
            ax1.legend(loc='best')
            ax1.grid(True, alpha=0.3)
            ax1.xaxis.set_major_formatter(DateFormatter("%Y-%m"))

            # 绘制成交量
            ax2.bar(data.index, data['Volume'],
                    color=np.where(data['Volume'] > data['Volume_MA'],
                                'green', 'gray'),
                    alpha=0.6, width=0.6)
            ax2.plot(data.index, data['Volume_MA'],
                    color='red', linewidth=1, label='Volume MA(5)')
            ax2.set_ylabel('Volume')
            ax2.legend(loc='best')
            ax2.grid(True, alpha=0.3)


            # 调整布局
            self.price_chart.figure.tight_layout()
            self.price_chart.draw()


            if len(data) > 1000:
                # 使用简单抽样减少数据点
                step = len(data) // 500
                plot_data = data.iloc[::step].copy()
                plot_signals = result.signals.loc[data.index].iloc[::step].copy()
            else:
                plot_data = data
                plot_signals = result.signals.loc[data.index]

        finally:
            QApplication.restoreOverrideCursor()

    def _plot_equity_curve(self, positions):
        """绘制资金曲线（支持时间范围）"""
        self.equity_chart.figure.clear()
        ax = self.equity_chart.figure.add_subplot(111)

        # 筛选时间范围
        if self.start_date and self.end_date:
            mask = (positions.index >= self.start_date) & (positions.index <= self.end_date)
            positions = positions.loc[mask]

        ax.plot(positions.index, positions['total'], label='Equity')
        ax.set_title("资金曲线")
        self.equity_chart.draw()

    def _show_metrics(self, metrics):
        """显示绩效指标"""
        text = "<b>回测结果:</b><br>"
        for k, v in metrics.items():
            text += f"{k}: {v}<br>"
        self.metrics_label.setText(text)


class MplCanvas(FigureCanvas):
    """支持中文的Matplotlib画布"""

    def __init__(self, parent=None, width=10, height=4, dpi=100):
        # 字体设置已在 gui.charts 中统一配置，这里不再重复设置
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)

    def clear(self):
        self.fig.clear()

    def plot_equity(self, dates, values):
        """绘制资金曲线（支持中文）"""
        self.clear()
        ax = self.fig.add_subplot(111)

        ax.plot(dates, values, label='资金曲线')
        ax.set_title("回测结果", fontsize=14)
        ax.set_xlabel("日期")
        ax.set_ylabel("资产")
        ax.legend()

        self.draw()

